from fastapi import FastAPI, HTTPException, UploadFile, File, Form, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
import asyncio
import json
import os
import uuid
import base64
import requests
import time
from datetime import datetime
import logging
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

app = FastAPI(title="Bailian Chat API")

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 开发环境允许所有来源，生产环境应限制
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 模拟数据存储
conversations = []
messages = {}

# 模型定义
class Message(BaseModel):
    conversation_id: Optional[str] = None
    message: str
    model: Optional[str] = "qwen2.5-7b"
    attachments: Optional[List[Dict[str, str]]] = []

class Conversation(BaseModel):
    title: Optional[str] = "新对话"

# 阿里云百炼API调用
def call_bailian_api(message, model, attachments=None):
    """调用阿里云百炼API"""
    try:
        # 获取API密钥和端点
        api_key = os.getenv("BAILIAN_API_KEY")
        if not api_key:
            logger.error("Missing BAILIAN_API_KEY in environment variables")
            return f"API密钥未配置，无法访问阿里云百炼服务。收到的消息：\"{message}\"。"
        
        # 阿里云通义千问API端点 - 确保使用正确的URL
        api_url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
        
        # 准备请求头
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        
        # 准备消息内容
        input_content = message
        
        # 如果有图像附件，构建多模态输入
        if attachments and any(a.get("type") == "image" for a in attachments):
            input_content = []
            # 首先添加文本内容
            input_content.append({"text": message})
            
            # 添加图像内容
            for attachment in attachments:
                if attachment.get("type") == "image":
                    image_url = attachment.get("url", "")
                    # 如果是base64数据
                    if image_url.startswith("data:image"):
                        # 从base64 URL中提取图像数据
                        image_data = image_url.split(",", 1)[1] if "," in image_url else ""
                        input_content.append({"image": image_data})
                    else:
                        # 如果是URL
                        input_content.append({"image": image_url})
        
        # 构建请求体
        # 根据模型名称确定正确的模型ID
        model_id = model.lower()
        
        # 确保模型ID符合阿里云百炼API要求
        if model_id.startswith("qwen"):
            # 通义千问模型命名格式：qwen-xxx
            if "omni" in model_id:
                model_id = "qwen-max"  # 修正为官方支持的模型ID
            elif "2.5-7b" in model_id:
                model_id = "qwen-max"  # 使用qwen-max作为默认
            else:
                model_id = "qwen-max"  # 如果不确定，使用qwen-max
        elif model_id.startswith("llama"):
            model_id = "llama3-8b-instruct"  # 使用支持的Llama模型
        else:
            model_id = "qwen-max"  # 默认使用qwen-max
            
        request_body = {
            "model": model_id,
            "input": {
                "messages": [
                    {
                        "role": "user",
                        "content": input_content
                    }
                ]
            },
            "parameters": {
                "result_format": "text",
                "seed": 1234,
                "temperature": 0.7,
                "top_p": 0.8
            }
        }
        
        logger.info(f"Sending request to Bailian API with model: {model_id}")
        response = requests.post(api_url, headers=headers, json=request_body)
        
        # 检查响应状态
        if response.status_code != 200:
            logger.error(f"Bailian API request failed: {response.status_code} - {response.text}")
            return f"API请求失败，状态码：{response.status_code}。收到的消息：\"{message}\"。"
        
        # 解析响应
        data = response.json()
        logger.info(f"Response from Bailian API: {json.dumps(data)[:200]}...")
        
        # 从响应中提取输出内容
        if "output" in data and "text" in data["output"]:
            return data["output"]["text"]
        else:
            logger.error(f"Unexpected response format: {data}")
            return f"收到意外的响应格式。收到的消息：\"{message}\"。"
            
    except Exception as e:
        logger.exception(f"Error calling Bailian API: {str(e)}")
        return f"调用API时发生错误：{str(e)}。收到的消息：\"{message}\"。"

# 直接调用API，不使用模拟模式
async def call_api_with_streaming(message, model, attachments=None):
    """直接调用通义千问API，支持流式输出"""
    try:
        # 获取API密钥
        api_key = os.getenv("BAILIAN_API_KEY")
        if not api_key:
            logger.error("Missing API key in environment variables")
            yield f"data: {{\"text\":\"API密钥未配置，请在后端.env文件中设置BAILIAN_API_KEY\"}}\n\n"
            yield "event: done\ndata: done\n\n"
            return
        
        # 通义千问API端点 - 为大模型设置正确的URL
        api_url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
        
        # 准备请求头
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}",
            "Accept": "text/event-stream",
            "X-DashScope-SSE": "enable"  # 确保启用SSE
        }
        
        # 准备消息内容
        input_content = message
        
        # 如果有图像附件，构建多模态输入
        if attachments and any(a.get("type") == "image" for a in attachments):
            input_content = []
            # 首先添加文本内容
            input_content.append({"text": message})
            
            # 添加图像内容
            for attachment in attachments:
                if attachment.get("type") == "image":
                    image_url = attachment.get("url", "")
                    # 如果是base64数据
                    if image_url.startswith("data:image"):
                        # 从base64 URL中提取图像数据
                        image_data = image_url.split(",", 1)[1] if "," in image_url else ""
                        input_content.append({"image": image_data})
                    else:
                        # 如果是URL
                        input_content.append({"image": image_url})
        
        # 根据模型名称确定正确的模型ID
        model_id = model.lower()
        
        # 确保模型ID符合阿里云百炼API要求
        if model_id.startswith("qwen"):
            # 通义千问模型命名格式：qwen-xxx
            if "omni" in model_id:
                model_id = "qwen-max"  # 修正为官方支持的模型ID
            elif "2.5-7b" in model_id:
                model_id = "qwen-max"  # 使用qwen-max作为默认
            else:
                model_id = "qwen-max"  # 如果不确定，使用qwen-max
        elif model_id.startswith("llama"):
            model_id = "llama3-8b-instruct"  # 使用支持的Llama模型
        else:
            model_id = "qwen-max"  # 默认使用qwen-max
        
        # 构建请求体    
        request_body = {
            "model": model_id,
            "input": {
                "messages": [
                    {
                        "role": "user",
                        "content": input_content
                    }
                ]
            },
            "parameters": {
                "result_format": "text",
                "incremental_output": True,  # 启用增量输出
                "seed": 1234,
                "temperature": 0.7,
                "top_p": 0.8
            }
        }
        
        logger.info(f"Sending streaming request to Bailian API with model: {model_id}")
        logger.info(f"Request body: {json.dumps(request_body)[:500]}...")
        
        # 尝试非流式方式获取响应（临时解决方案）
        try:
            # 先尝试非流式请求（确保API工作正常）
            non_streaming_headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {api_key}"
            }
            
            # 使用非增量输出
            non_streaming_body = request_body.copy()
            non_streaming_body["parameters"]["incremental_output"] = False
            
            logger.info("Attempting non-streaming request first...")
            response = requests.post(api_url, headers=non_streaming_headers, json=non_streaming_body, timeout=30)
            
            if response.status_code != 200:
                error_message = f"API请求失败，状态码：{response.status_code} - {response.text}"
                logger.error(error_message)
                yield f"data: {{\"text\":\"API请求失败: {response.status_code}\"}}\n\n"
                yield "event: done\ndata: done\n\n"
                return
            
            # 解析响应
            data = response.json()
            logger.info(f"Non-streaming response received: {json.dumps(data)[:200]}...")
            
            if "output" in data and "text" in data["output"]:
                # 模拟流式传输（分段发送）
                text = data["output"]["text"]
                # 分割文本为较小的块
                chunk_size = 20  # 每个块包含的字符数
                text_chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
                
                for chunk in text_chunks:
                    yield f"data: {{\"text\":\"{chunk}\"}}\n\n"
                    await asyncio.sleep(0.1)  # 小延迟模拟打字效果
                
                yield "event: done\ndata: done\n\n"
                return
            else:
                logger.error(f"Unexpected response format: {data}")
                yield f"data: {{\"text\":\"收到意外的响应格式\"}}\n\n"
                yield "event: done\ndata: done\n\n"
                return
                
        except Exception as e:
            logger.exception(f"Non-streaming request failed: {str(e)}")
            # 继续尝试流式请求
        
        # 发送请求并处理流式响应
        logger.info("Attempting streaming request...")
        with requests.post(api_url, headers=headers, json=request_body, stream=True, timeout=60) as response:
            if response.status_code != 200:
                error_message = f"流式API请求失败，状态码：{response.status_code}"
                logger.error(f"{error_message} - {response.text}")
                yield f"data: {{\"text\":\"{error_message}\"}}\n\n"
                yield "event: done\ndata: done\n\n"
                return
            
            # 记录响应头信息以便调试
            logger.info(f"Response headers: {response.headers}")
            
            # 处理流式响应
            full_response = ""
            for line in response.iter_lines():
                if line:
                    line_text = line.decode('utf-8')
                    logger.debug(f"Stream line received: {line_text}")
                    
                    if line_text.startswith("data:"):
                        try:
                            data = json.loads(line_text[5:].strip())
                            logger.debug(f"Stream data parsed: {data}")
                            
                            # 提取回复文本
                            if "output" in data and "text" in data["output"]:
                                text_chunk = data["output"]["text"]
                                full_response += text_chunk
                                logger.debug(f"Extracted text chunk: {text_chunk}")
                                
                                # 发送给客户端
                                yield f"data: {{\"text\":\"{text_chunk}\"}}\n\n"
                            else:
                                logger.warning(f"No text in chunk: {data}")
                        except json.JSONDecodeError as e:
                            logger.error(f"Failed to parse JSON: {e}, line: {line_text}")
            
            if not full_response:
                logger.warning("No response received from streaming API")
                yield f"data: {{\"text\":\"未能从API获取响应，请稍后再试\"}}\n\n"
            
            # 完成事件
            yield "event: done\ndata: done\n\n"
            
    except Exception as e:
        error_message = f"调用API时发生错误：{str(e)}"
        logger.exception(error_message)
        yield f"data: {{\"text\":\"{error_message}\"}}\n\n"
        yield "event: done\ndata: done\n\n"

# API路由
@app.get("/conversations")
async def get_conversations():
    if not conversations:
        # 如果没有会话，创建一个初始会话
        new_id = f"conv{uuid.uuid4().hex[:8]}"
        new_conv = {
            "id": new_id,
            "title": "新对话",
            "created_at": datetime.now().isoformat()
        }
        conversations.append(new_conv)
        messages[new_id] = []
    return conversations

@app.post("/conversations")
async def create_conversation(conversation: Conversation):
    new_id = f"conv{uuid.uuid4().hex[:8]}"
    new_conv = {
        "id": new_id,
        "title": conversation.title,
        "created_at": datetime.now().isoformat()
    }
    conversations.append(new_conv)
    messages[new_id] = []
    return new_conv

@app.get("/conversations/{conversation_id}/messages")
async def get_messages(conversation_id: str):
    if conversation_id not in messages:
        raise HTTPException(status_code=404, detail="Conversation not found")
    return messages.get(conversation_id, [])

@app.post("/chat")
async def send_message(message: Message):
    try:
        # 获取或创建会话
        conversation_id = message.conversation_id
        if not conversation_id:
            new_id = f"conv{uuid.uuid4().hex[:8]}"
            conversations.append({
                "id": new_id,
                "title": "新对话",
                "created_at": datetime.now().isoformat()
            })
            messages[new_id] = []
            conversation_id = new_id
        elif conversation_id not in messages:
            raise HTTPException(status_code=404, detail="Conversation not found")
        
        # 添加用户消息
        user_msg_id = f"msg{uuid.uuid4().hex[:8]}"
        user_message = {
            "id": user_msg_id,
            "role": "user",
            "content": message.message,
            "attachments": message.attachments,
            "timestamp": datetime.now().isoformat()
        }
        messages[conversation_id].append(user_message)
        
        # 直接调用通义千问API
        response_text = call_bailian_api(message.message, message.model, message.attachments)
        
        # 处理响应
        ai_msg_id = f"msg{uuid.uuid4().hex[:8]}"
        ai_message = {
            "id": ai_msg_id,
            "role": "assistant",
            "content": response_text,
            "timestamp": datetime.now().isoformat()
        }
        messages[conversation_id].append(ai_message)
        
        # 如果是第一条消息，更新会话标题
        if len(messages[conversation_id]) == 2:  # 用户消息和AI回复
            for i, conv in enumerate(conversations):
                if conv["id"] == conversation_id:
                    # 使用用户消息的前15个字符作为标题
                    title = message.message[:15] + ("..." if len(message.message) > 15 else "")
                    conversations[i]["title"] = title
                    break
        
        return ai_message
        
    except Exception as e:
        logger.error(f"Error in send_message: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/stream")
async def stream_response(request: Request):
    """流式响应接口"""
    try:
        # 从查询参数获取数据
        query_params = request.query_params
        data_str = query_params.get("data", "{}")
        data = json.loads(data_str)
        
        # 提取消息内容
        message = data.get("message", "")
        model = data.get("model", "qwen2.5-7b")
        attachments = data.get("attachments", [])
        conversation_id = data.get("conversation_id")
        
        # 创建或获取会话
        if not conversation_id:
            new_id = f"conv{uuid.uuid4().hex[:8]}"
            conversations.append({
                "id": new_id,
                "title": "新对话",
                "created_at": datetime.now().isoformat()
            })
            messages[new_id] = []
            conversation_id = new_id
        elif conversation_id not in messages:
            return StreamingResponse(
                content=iter([f"data: {{\"text\":\"会话不存在\"}}\n\n".encode('utf-8'), b"event: done\ndata: done\n\n"]),
                media_type="text/event-stream"
            )
        
        # 添加用户消息
        user_msg_id = f"msg{uuid.uuid4().hex[:8]}"
        user_message = {
            "id": user_msg_id,
            "role": "user",
            "content": message,
            "attachments": attachments,
            "timestamp": datetime.now().isoformat()
        }
        messages[conversation_id].append(user_message)
        
        # 如果是会话的第一条消息，设置会话标题
        if len(messages[conversation_id]) == 1:
            for i, conv in enumerate(conversations):
                if conv["id"] == conversation_id:
                    title = message[:15] + ("..." if len(message) > 15 else "")
                    conversations[i]["title"] = title
                    break
        
        # 使用流式API返回响应
        return StreamingResponse(
            content=call_api_with_streaming(message, model, attachments),
            media_type="text/event-stream"
        )
        
    except Exception as e:
        logger.exception(f"Error in stream_response: {str(e)}")
        error_message = f"处理请求时出错：{str(e)}"
        return StreamingResponse(
            content=iter([f"data: {{\"text\":\"{error_message}\"}}\n\n".encode(), b"event: done\ndata: done\n\n"]),
            media_type="text/event-stream"
        )

@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    """上传文件API"""
    try:
        # 生成唯一文件名
        filename = f"{uuid.uuid4().hex}_{file.filename}"
        upload_dir = "uploads"
        
        # 确保上传目录存在
        os.makedirs(upload_dir, exist_ok=True)
        
        # 保存文件
        file_path = os.path.join(upload_dir, filename)
        with open(file_path, "wb") as f:
            f.write(await file.read())
        
        # 如果是图像文件，转换为base64
        if file.content_type.startswith("image/"):
            with open(file_path, "rb") as f:
                file_content = f.read()
                base64_data = base64.b64encode(file_content).decode("utf-8")
                data_url = f"data:{file.content_type};base64,{base64_data}"
                
                return {
                    "url": data_url,
                    "type": "image",
                    "name": file.filename
                }
        
        # 其他类型文件返回文件路径
        return {
            "url": f"/uploads/{filename}",
            "type": "file",
            "name": file.filename
        }
            
    except Exception as e:
        logger.exception(f"Error uploading file: {str(e)}")
        raise HTTPException(status_code=500, detail=f"文件上传失败: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=True)